{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Copyright 2025 Google LLC\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Virtual Try-On - Generation at Scale\n",
        "\n",
        "<table align=\"left\">\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-creative-studio/blob/main/experiments/Imagen_Product_Recontext/imagen_product_recontext_at_scale.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://www.gstatic.com/pantheon/images/bigquery/welcome_page/colab-logo.svg\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%vertex-ai-creative-studio%2Fmain%2Fexperiments%2FImagen_Product_Recontext%imagen_product_recontext_at_scale.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-creative-studio/main/experiments/Imagen_Product_Recontext/imagen_product_recontext_at_scale.ipynb\">\n",
        "      <img src=\"https://www.gstatic.com/images/branding/gcpiconscolors/vertexai/v1/32px.svg\" alt=\"Vertex AI logo\"><br> Open in Vertex AI Workbench\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/bigquery/import?url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-creative-studio/main/experiments/Imagen_Product_Recontext/imagen_product_recontext_at_scale.ipynb\">\n",
        "      <img src=\"https://www.gstatic.com/images/branding/gcpiconscolors/bigquery/v1/32px.svg\" alt=\"BigQuery Studio logo\"><br> Open in BigQuery Studio\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-creative-studio/blob/main/experiments/Imagen_Product_Recontext/imagen_product_recontext_at_scale.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://upload.wikimedia.org/wikipedia/commons/9/91/Octicons-mark-github.svg\" alt=\"GitHub logo\"><br> View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "</table>\n",
        "\n",
        "<div style=\"clear: both;\"></div>"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "YtECZUtNtuvZ",
      "metadata": {
        "id": "YtECZUtNtuvZ"
      },
      "source": [
        " Author : Layolin Jesudhass , Gen AI Solution Architect\n"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "WRyO_sF8uX8F",
      "metadata": {
        "id": "WRyO_sF8uX8F"
      },
      "source": [
        "# Virtual Try-On\n",
        " This notebook demonstrates the virtual try-on process.\n",
        " It includes the uploaded person image and try-on results for selected outfits.\n",
        " Try on multiple outfits in parallel and compare your looks side by side in a single view."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "cQNFLDFajJNW",
      "metadata": {
        "id": "cQNFLDFajJNW"
      },
      "outputs": [],
      "source": [
        "from PIL import Image\n",
        "from IPython.display import display\n",
        "import base64\n",
        "import io\n",
        "import time\n",
        "import os\n",
        "import concurrent.futures\n",
        "from IPython.display import display\n",
        "from google.cloud import aiplatform\n",
        "from google.cloud.aiplatform.gapic import PredictionServiceClient\n",
        "from IPython.display import display, HTML\n",
        "\n",
        "# Constants\n",
        "PROJECT_ID = \"consumer-genai-experiments\"\n",
        "LOCATION = \"us-central1\"\n",
        "MODEL_ID = \"virtual-try-on-exp-05-31\"\n",
        "PRODUCT_IMAGE_FILES = [\"red.jpg\", \"green.png\", \"dress.png\", \"blue.png\", \"yellow.png\"]\n",
        "TARGET_SIZE = (250, 550)\n",
        "model_endpoint = f\"projects/{PROJECT_ID}/locations/{LOCATION}/publishers/google/models/{MODEL_ID}\"\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "oo4GF0hMlsPu",
      "metadata": {
        "id": "oo4GF0hMlsPu"
      },
      "outputs": [],
      "source": [
        "# --- Utility functions ---\n",
        "\n",
        "def load_image_bytes(path):\n",
        "    with open(path, \"rb\") as f:\n",
        "        return f.read()\n",
        "\n",
        "def encode_image(img_bytes):\n",
        "    return base64.b64encode(img_bytes).decode(\"utf-8\")\n",
        "\n",
        "def prediction_to_pil_image(prediction, size=TARGET_SIZE):\n",
        "    encoded = prediction[\"bytesBase64Encoded\"]\n",
        "    decoded = base64.b64decode(encoded)\n",
        "    image = Image.open(io.BytesIO(decoded)).convert(\"RGB\")\n",
        "    return image.resize(size)\n",
        "\n",
        "def run_tryon(person_b64, name, b64):\n",
        "    start = time.time()\n",
        "    client = PredictionServiceClient(client_options={\"api_endpoint\": f\"{LOCATION}-aiplatform.googleapis.com\"})\n",
        "    instances = [{\n",
        "        \"personImage\": {\"image\": {\"bytesBase64Encoded\": person_b64}},\n",
        "        \"productImages\": [{\"image\": {\"bytesBase64Encoded\": b64}}],\n",
        "    }]\n",
        "    response = client.predict(endpoint=model_endpoint, instances=instances, parameters={})\n",
        "    elapsed = time.time() - start\n",
        "    output_img = prediction_to_pil_image(response.predictions[0])\n",
        "    return output_img, elapsed\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "nN3GpWRdjJQL",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 567
        },
        "id": "nN3GpWRdjJQL",
        "outputId": "5279e3a2-3017-4792-d9e2-389ae1ae77c9"
      },
      "outputs": [],
      "source": [
        "# Load and encode person image\n",
        "person_image_path = \"model.png\"\n",
        "with open(person_image_path, \"rb\") as f:\n",
        "    person_bytes = f.read()\n",
        "\n",
        "def encode_image(img_bytes):\n",
        "    return base64.b64encode(img_bytes).decode(\"utf-8\")\n",
        "\n",
        "person_b64 = encode_image(person_bytes)\n",
        "\n",
        "# Open and resize the person image for display (max 250x550, keeping aspect ratio)\n",
        "person_img = Image.open(person_image_path)\n",
        "person_img.thumbnail((250, 550))  # Resize in-place maintaining aspect ratio\n",
        "\n",
        "# Display inline small\n",
        "display(person_img)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "-epxz4XVjJT1",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 289
        },
        "id": "-epxz4XVjJT1",
        "outputId": "b41b7460-322f-40df-ec91-6fef6eec7e69"
      },
      "outputs": [],
      "source": [
        "def pil_image_to_base64(img):\n",
        "    buf = io.BytesIO()\n",
        "    img.save(buf, format=\"PNG\")\n",
        "    return base64.b64encode(buf.getvalue()).decode()\n",
        "\n",
        "product_pil_images = []\n",
        "for file_name in PRODUCT_IMAGE_FILES:\n",
        "    img = Image.open(file_name)\n",
        "    img = img.resize((150, 250))  # fixed size (width=150, height=200)\n",
        "    product_pil_images.append((file_name, img))\n",
        "\n",
        "imgs_html = \"\"\n",
        "for name, img in product_pil_images:\n",
        "    img_b64 = pil_image_to_base64(img)\n",
        "    imgs_html += f\"\"\"\n",
        "    <div style=\"text-align:center; margin-right:15px;\">\n",
        "        <img src=\"data:image/png;base64,{img_b64}\" style=\"max-height:250px; max-width:150px;\"><br>\n",
        "        <b>{name}</b><br>\n",
        "    </div>\n",
        "    \"\"\"\n",
        "\n",
        "display(HTML(f\"<div style='display:flex; align-items:center;'>{imgs_html}</div>\"))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "xOOUBthF5-px",
      "metadata": {
        "id": "xOOUBthF5-px"
      },
      "outputs": [],
      "source": [
        "# Prepare product_data: a list of (file_name, base64_encoded_image) tuples\n",
        "product_data = []\n",
        "for file_name in PRODUCT_IMAGE_FILES:\n",
        "    with open(file_name, \"rb\") as f:\n",
        "        img_bytes = f.read()\n",
        "        img_b64 = base64.b64encode(img_bytes).decode(\"utf-8\")\n",
        "        product_data.append((file_name, img_b64))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "ij5MujBLnOSc",
      "metadata": {
        "id": "ij5MujBLnOSc"
      },
      "outputs": [],
      "source": [
        "# --- Run inference in parallel ---\n",
        "results = []\n",
        "\n",
        "def run_thread(item):\n",
        "    name, b64 = item\n",
        "    out_img, elapsed = run_tryon(person_b64, name, b64)\n",
        "    return (name, out_img, elapsed)\n",
        "\n",
        "with concurrent.futures.ThreadPoolExecutor() as executor:\n",
        "    futures = [executor.submit(run_thread, pd) for pd in product_data]\n",
        "    for f in concurrent.futures.as_completed(futures):\n",
        "        results.append(f.result())"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "YSrUFpx_nOVP",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 306
        },
        "id": "YSrUFpx_nOVP",
        "outputId": "9c925742-575a-40a0-ad6b-0e521bce0c03"
      },
      "outputs": [],
      "source": [
        "ordered_results = sorted(results, key=lambda x: PRODUCT_IMAGE_FILES.index(x[0]))\n",
        "\n",
        "imgs_html = \"\"\n",
        "for name, out_img, elapsed in ordered_results:\n",
        "    img_b64 = pil_image_to_base64(out_img)\n",
        "    imgs_html += f\"\"\"\n",
        "    <div style=\"text-align:center; margin-right:15px;\">\n",
        "        <img src=\"data:image/png;base64,{img_b64}\" style=\"max-height:250px; max-width:150px;\"><br>\n",
        "        <b>{\"Time Taken\"}</b><br>\n",
        "        <small>{elapsed:.2f}s</small>\n",
        "    </div>\n",
        "    \"\"\"\n",
        "\n",
        "display(HTML(f\"<div style='display:flex; align-items:center;'>{imgs_html}</div>\"))"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "VTO at scale",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.10.10"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}
